from paddleocr import PaddleOCR, PPStructure
import fitz
from docx import Document
from docx.shared import Inches
import os
import numpy as np
from PIL import Image

def pdf_to_docx_with_structure(pdf_path, docx_path):
    # 初始化 PaddleOCR 和 PPStructure
    ocr = PaddleOCR(use_angle_cls=True, lang='ch')
    table_engine = PPStructure(table=True, ocr=False, show_log=False)
    
    # 创建 Word 文档
    doc = Document()
    
    # 打开 PDF 文件
    pdf = fitz.open(pdf_path)
    
    for page_num in range(len(pdf)):
        page = pdf[page_num]
        doc.add_heading(f'第 {page_num + 1} 页', level=1)
        
        # 将 PDF 页面转换为图片
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 提高分辨率
        img_path = f"temp_page_{page_num}.png"
        pix.save(img_path)
        
        # 使用结构化识别
        img = Image.open(img_path)
        result = table_engine(np.array(img))
        
        # 处理识别结果
        for region in result:
            if region['type'] == 'table':  # 处理表格
                cells = region['res']['cells']
                if not cells:
                    continue
                    
                # 获取表格行列数
                rows = max(cell[0][1] for cell in cells) + 1
                cols = max(cell[0][0] for cell in cells) + 1
                
                # 创建 Word 表格
                table = doc.add_table(rows=rows, cols=cols)
                table.style = 'Table Grid'
                
                # 填充表格内容
                for cell in cells:
                    col, row = cell[0]
                    text = cell[1]
                    table.cell(row, col).text = text
                
                # 调整表格格式
                adjust_table_format(table)
                
                # 表格后添加空行
                doc.add_paragraph()
                
            else:  # 处理普通文本
                # 使用标准 OCR 识别
                ocr_result = ocr.ocr(np.array(img), cls=True)
                if ocr_result:
                    for line in ocr_result[0]:
                        text = line[1][0]
                        doc.add_paragraph(text)
        
        # 删除临时文件
        os.remove(img_path)
        
        # 添加分页符（除了最后一页）
        if page_num < len(pdf) - 1:
            doc.add_page_break()
    
    # 保存文档
    doc.save(docx_path)
    pdf.close()

def adjust_table_format(table):
    """调整表格格式"""
    from docx.shared import Pt
    from docx.enum.text import WD_ALIGN_PARAGRAPH
    
    # 设置表格格式
    for row in table.rows:
        for cell in row.cells:
            # 设置单元格内文字格式
            paragraph = cell.paragraphs[0]
            paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            # 设置字体
            run = paragraph.runs[0] if paragraph.runs else paragraph.add_run()
            run.font.size = Pt(10)
            
            # 设置单元格垂直对齐
            cell.vertical_alignment = 1  # 1 表示居中对齐

# 使用示例
if __name__ == "__main__":
    pdf_path = "pdfFiles/GB 14193-2009 液化气体气瓶充装规定 .pdf"
    docx_path = "output.docx"
    pdf_to_docx_with_structure(pdf_path, docx_path)